import cv2
import argparse
import os

import sys
from pathlib import Path
 
sys.path.append(str(Path(__file__).resolve().parent.parent))

from file_util.file_util import GetAllSubFile
import xml.etree.ElementTree as ET



def draw_image_with_anno(new_img_file,new_anno_file):
		try:
			tree = ET.parse(new_anno_file)
			root = tree.getroot() 
		except:
			print(new_anno_file + " is not exist.")
			return
		img = cv2.imread(new_img_file)
		for object in root:
			if object.tag == "object":
				bndbox = object[4][0]
				xmin = int(object[4][0].text)
				ymin = int(object[4][1].text)
				xmax= int(object[4][2].text)
				ymax = int(object[4][3].text)
				cv2.rectangle(img,(xmin,ymin),(xmax,ymax),(0,0,255),5)
		cv2.imwrite(new_img_file,img)
def update_anno(anno_file, w, h, new_anno_file, flip_flag):
		try:
			tree = ET.parse(anno_file)
			root = tree.getroot() 
		except:
			print(anno_file + " is not exist.")
			return
		remove_objects = []
		for object in root:
			if object.tag == "object":
				bndbox = object[4][0]
				xmin = int(object[4][0].text)
				ymin = int(object[4][1].text)
				xmax= int(object[4][2].text)
				ymax = int(object[4][3].text)
				if flip_flag == 0: #v flip
					new_xmin = xmin
					new_xmax = xmax
					new_ymin = h - ymax
					new_ymax = h - ymin
				elif flip_flag == 1: #h flip
					new_xmin = w - xmax
					new_xmax = w - xmin
					new_ymin = ymin
					new_ymax = ymax
				elif flip_flag == -1: #both flip
					new_xmin = w - xmax
					new_xmax = w - xmin
					new_ymin = h - ymax
					new_ymax = h - ymin
				else:
					print("not support flip_flag.\n")
					return
				if (new_xmax < 0) or (new_ymax < 0):
					print(object[0].text)
					#root.remove(object)
					remove_objects.append(object)
					continue
				if (new_xmin > w) or (new_ymin > h):
					print(object[0].text)
					#root.remove(object)
					remove_objects.append(object)
					continue
				if new_xmin < 0:
					new_xmin = 0
				if new_ymin < 0 :
					new_ymin = 0
				if new_xmax > w :
					new_xmax = w
				if new_ymax > h :
					new_ymax = h
				object[4][0].text = str(new_xmin)
				object[4][1].text = str(new_ymin)
				object[4][2].text = str(new_xmax)
				object[4][3].text = str(new_ymax)
		for remove_object in remove_objects:
			root.remove(remove_object)
		tree.write(new_anno_file,encoding="utf-8", xml_declaration=True)
		return 
def data_flip(opt):
	img_files = GetAllSubFile(opt.root_data)
	for img_file in img_files:
		img = cv2.imread(img_file)
		h = img.shape[0]
		w = img.shape[1]
		path, img_file_name = os.path.split(img_file)
		img_name = img_file_name.rsplit('.',1)[0]
		img_flip_v = cv2.flip(img,0)
		img_flip_h = cv2.flip(img,1)
		img_flip_both = cv2.flip(img,-1)
		#print(opt.anon_dir)
		#print(opt.root_data);
		#print(img_name)
		#anno_file = os.path.join(opt.anon_dir,img_name + '.xml')
		print(len(opt.anon_dir),len(img_name))
		anno_file = opt.anon_dir + img_name + '.xml'
		print("anno_file=",anno_file,"opt.anon_dir=",opt.anon_dir)
		
		new_img_name = img_name+"_flip_v"
		new_anno_file = os.path.join(opt.anon_dir,new_img_name + '.xml')
		new_img_file = os.path.join(path,new_img_name +".jpg")
		cv2.imwrite(new_img_file,img_flip_v)
		update_anno(anno_file,w,h,new_anno_file,0)
		#draw_image_with_anno(new_img_file,new_anno_file);
		
		new_img_name = img_name+"_flip_h"
		new_anno_file = os.path.join(opt.anon_dir,new_img_name + '.xml')
		new_img_file = os.path.join(path,new_img_name +".jpg")
		cv2.imwrite(new_img_file,img_flip_h)
		update_anno(anno_file,w,h,new_anno_file,1)
		#draw_image_with_anno(new_img_file,new_anno_file);

		new_img_name = img_name+"_flip_both"
		new_anno_file = os.path.join(opt.anon_dir,new_img_name + '.xml')
		new_img_file = os.path.join(path,new_img_name +".jpg")
		cv2.imwrite(new_img_file,img_flip_both)
		update_anno(anno_file,w,h,new_anno_file,-1)	
		#draw_image_with_anno(new_img_file,new_anno_file);

if __name__ == '__main__':
   parser = argparse.ArgumentParser("")
   parser.add_argument('--root_data', type=str, default='./data', help='root data dir.')
   parser.add_argument('--anon_dir', type=str, default='./anno', help='anno dir.')
   opt = parser.parse_args()
   opt.anon_dir = opt.anon_dir.strip()
   #print(opt.anon_dir)
   #print(opt.root_data);
   #3
   data_flip(opt)